# -*- coding:utf8 -*-
# @Time : 2023/3/13 18:17
# @Author : WanJie Wu

from dataset import get_entity, gen_classify_entity


class SeqEntityMetric(object):
    def __init__(self, markup, processor):
        self.markup = markup
        self.processor = processor
        self._init()

    def _init(self):
        self.metric_detail = []
        self.loss = 0.
        self.steps = 0

    @staticmethod
    def calculate_indicator_by_dict(val):
        if val["pred_num"] == val["gold_num"] == val["correct_num"]:
            precision = 1
            recall = 1
            f1 = 1
        else:
            precision = round(val["correct_num"] / val["pred_num"], 3) if val["pred_num"] > 0 else 0.0
            recall = round(val["correct_num"] / val["gold_num"], 3) if val["gold_num"] > 0 else 0.0
            f1 = round(2 * precision * recall / (precision + recall), 3) if (precision + recall) > 0 else 0.0
        return {
            "precision": precision,
            "recall": recall,
            "f1": f1
        }

    def result(self):
        class_info = dict()
        t_gold_num, t_pred_num, t_correct_count = 0, 0, 0
        for detail in self.metric_detail:
            for key, per_metric in detail.items():
                if key in class_info:
                    class_info[key]["gold_num"] += per_metric["gold_num"]
                    class_info[key]["pred_num"] += per_metric["pred_num"]
                    class_info[key]["correct_num"] += per_metric["correct_num"]
                else:
                    class_info[key] = per_metric
                t_gold_num += per_metric["gold_num"]
                t_pred_num += per_metric["pred_num"]
                t_correct_count += per_metric["correct_num"]
        tmp_result = dict()
        for key, val in class_info.items():
            tmp_result[key] = self.calculate_indicator_by_dict(val)
        return {
            "loss": round(self.loss / self.steps, 3),
            "total": self.calculate_indicator_by_dict({"gold_num": t_gold_num, "pred_num": t_pred_num, "correct_num": t_correct_count}),
            "detail": tmp_result
        }

    @staticmethod
    def _calculate_metrics(gold_entities, predict_entities):
        gold_detail = gen_classify_entity(gold_entities)
        pred_detail = gen_classify_entity(predict_entities)

        detail = dict()
        for key, val in gold_detail.items():
            gold_num = len(val)
            if key not in pred_detail.keys():
                detail[key] = {
                    "gold_num": gold_num,
                    "pred_num": 0,
                    "correct_num": 0
                }
                continue
            pred_num = len(pred_detail[key])
            correct_count = len([pred for pred in pred_detail[key] if pred in val])
            detail[key] = {
                "gold_num": gold_num,
                "pred_num": pred_num,
                "correct_num": correct_count
            }
        return detail

    def update(self, gold_paths, pred_paths, loss):
        """"""
        self.loss += loss
        for gold_path, pred_path in zip(gold_paths, pred_paths):
            # [['PER', 0, 1], ['LOC', 3, 3]]
            self.steps += 1
            gold_entities = get_entity(gold_path, self.processor.id2label, self.markup)
            pred_entities = get_entity(pred_path, self.processor.id2label, self.markup)
            per_metric = self._calculate_metrics(gold_entities, pred_entities)
            self.metric_detail.append(per_metric)